
module YUVPicture (
    ColorMode( YUV420, YUV422, YUV444 ),
    SizeType, PixelRow, ColorPlane, PicInfo, Picture, Image,
    scaleWidth, scaleHeight, scaleSize,
    toYUV444,
    mbPixelCount, mbSize, picPixelCount, picSize,
    packPlane, unpackPlane, packPic, unpackPic, cropPlane, cropPic,
    PixelYUV,
    chop,
    hGetPic, getPic, hGetPics, getPics
  ) where


import qualified Data.ByteString.Lazy as B
import Data.Word
import Data.Int
import List
import Maybe
import System.IO

{-
-- Test data generation script
let pack = Data.ByteString.pack
let unpack = Data.ByteString.unpack
let s1 = replicate 4 [0..3]
let s2 = Data.List.transpose s1
let y1 = map (pack . map (toEnum . fromInteger)) s1
let y2 = map (pack . map (toEnum . fromInteger)) s2
let u1 = y1
let u2 = y2
let v1 = y1
let v2 = y2
let pic1 = (YUV420,[y1,u1,v1])
let pic2 = (YUV420,[y2,u2,v2])
-}


data ColorMode = YUV420 | YUV422 | YUV444
  deriving (Enum, Eq, Show, Read)

type SizeType = Int

type PixelRow = B.ByteString

type ColorPlane = [PixelRow]

type Picture = (ColorMode, [ColorPlane])

type PicInfo = (ColorMode, SizeType, SizeType)

type Image = (PicInfo, B.ByteString)


scaleWidth :: ColorMode -> (Rational,Rational)
scaleWidth m = case m of
  YUV420 -> (1/2,1/2)
  YUV422 -> (1/2,1/2)
  YUV444 -> (1,1)

scaleHeight :: ColorMode -> (Rational,Rational)
scaleHeight m = case m of
  YUV420 -> (1/2,1/2)
  YUV422 -> (1,1)
  YUV444 -> (1,1)

scaleSize :: ColorMode -> (Rational,Rational)
scaleSize m = case m of
  YUV420 -> (1/4,1/4)
  YUV422 -> (1/2,1/2)
  YUV444 -> (1,1)

mbPixelCount :: ColorMode -> Int
mbPixelCount _ = 16*16

mbSize :: ColorMode -> Int
mbSize m = round
  . (*) (toRational $ mbPixelCount m)
  . (1+) . uncurry (+) $ scaleSize m

picPixelCount :: PicInfo -> Int
picPixelCount (_,w,h) = w*h

picSize :: PicInfo -> Int
picSize i@(m,_,_) = round
  . (*) (toRational $ picPixelCount i)
  . (1+) . uncurry (+) $ scaleSize m


-- unpackPic :: Picture -> [Picture]
-- unpackPic (m,[y,u,v]) = zipWith4 mkPic (repeat m) ys us vs
--   where
--     mkPic m y u v = (m,[y,u,v])
--     (rWCb,rWCr) = scaleWidth m
--     (rHCb,rHCr) = scaleHeight m
--     scale sz = round . (*) (toRational sz)
--     cWY = 16
--     cHY = 16
--     cWCb = scale cWY rWCb
--     cHCb = scale cHY rHCb
--     cWCr = scale cWY rWCr
--     cHCr = scale cHY rHCr
--     unpackMB :: Int -> Int -> ColorPlane -> [ColorPlane]
--     unpackMB w h = concat . map (map (map B.pack) . transpose) . chop h . map (chop w . B.unpack)
--     ys = unpackMB cWY  cHY  y
--     us = unpackMB cWCb cHCb u
--     vs = unpackMB cWCr cHCr v

{-
a b c d e f
g h i j k l
m n o p q r
s t u v w x

map (chop 2)

a b, c d, e f
g h, i j, k l
m n, o p, q r
s t, u v, w x

chop 2

a b, c d, e f
g h, i j, k l

m n, o p, q r
s t, u v, w x

transpose

a b, c d, e f
m n, o p, q r

g h, i j, k l
s t, u v, w x

concat

a b, c d, e f
m n, o p, q r
g h, i j, k l
s t, u v, w x

map concat

a b c d e f
m n o p q r
g h i j k l
s t u v w x
-}
chop      :: Int -> [a] -> [[a]]
chop _ [] = []
chop n ls = take n ls : chop n (drop n ls)
chopB      :: Int64 -> B.ByteString -> [B.ByteString]
chopB n bs | B.length bs == 0 = []
           | otherwise        = B.take n bs : chopB n (B.drop n bs)

packPlane :: Int -> Int -> PixelRow -> ColorPlane
packPlane _ 0 _ = []
packPlane w h s = B.take w_ s : packPlane w (h-1) (B.drop w_ s)
  where w_ = fromIntegral w

unpackPlane :: Int -> Int -> ColorPlane -> [[ColorPlane]]
unpackPlane w h = map transpose . chop h . map (chopB $ fromIntegral w)

packPic :: Image -> Picture
packPic (i@(m,w,h),s) = (m,[y,u,v])
  where
    (rSzCb,rSzCr) = scaleSize m
    (rWCb,rWCr) = scaleWidth m
    (rHCb,rHCr) = scaleHeight m
    szY  = fromIntegral $ picPixelCount i
    scale sz = round . (*) (toRational sz)
    szCb = fromIntegral $ scale szY rSzCb
    szCr = fromIntegral $ scale szY rSzCr
    y = packPlane  w              h             $ B.take szY                      $ s
    u = packPlane (scale w rWCb) (scale h rHCb) $ B.take szCb . B.drop szY        $ s
    v = packPlane (scale w rWCr) (scale h rHCr) $ B.take szCr . B.drop (szY+szCb) $ s

unpackPic :: Picture -> [[Picture]]
unpackPic (m,[y,u,v]) = zipWith3 (zipWith3 (\y u v -> (m,[y,u,v]))) ys us vs
  where
    (rWCb,rWCr) = scaleWidth m
    (rHCb,rHCr) = scaleHeight m
    scale sz = round . (*) (toRational sz)
    ys = unpackPlane  16              16             $ y
    us = unpackPlane (scale 16 rWCb) (scale 16 rHCb) $ u
    vs = unpackPlane (scale 16 rWCr) (scale 16 rHCr) $ v


cropPlane :: (Int,Int) -> (Int,Int) -> ColorPlane -> ColorPlane
cropPlane (x0,y0) (x1,y1)
  | all (>=0) [x0,y0] = map (cropB x0 x1 . flip B.append (linePad dx)) . crop y0 y1
  | y0 < 0            = cropPlane (x0,0) (x1,dy)
          . (++) (replicate (-y0) . B.replicate (fromIntegral dx) $ 128)
  | x0 < 0            = cropPlane (0,y0) (dx,y1)
          . map (B.append (linePad (-x0)))
  where
    linePad n = B.replicate (fromIntegral n) $ 128
    dx = x1 - x0
    dy = y1 - y0
    crop  l r =   drop l  .   take r
    cropB l r = B.drop l_ . B.take r_
      where l_ = fromIntegral l
            r_ = fromIntegral r

cropPic :: (Int,Int) -> (Int,Int) -> Picture -> Picture
cropPic lb@(xmin,ymin) ub@(xmax,ymax) (m,[y,u,v]) = (m,[ys,us,vs])
  where
    convRange sel bound (x,y) = (xx,yy)
      where xx = bound . (toRational x *) . sel $ scaleWidth m
            yy = bound . (toRational y *) . sel $ scaleHeight m
    ys = cropPlane lb                       ub                         y
    us = cropPlane (convRange fst floor lb) (convRange fst ceiling ub) u
    vs = cropPlane (convRange snd floor lb) (convRange snd ceiling ub) v


toYUV444 :: Picture -> Picture
toYUV444 (m,[y,u,v]) = (YUV444,[y_,u_,v_])
  where
    up  n ls =   concatMap (  replicate n) ls
    upB n bs = B.concatMap (B.replicate n) bs
    (rWCb,rWCr) = scaleWidth m
    (rHCb,rHCr) = scaleHeight m
    rc r = floor (1 / r)
    y_ = y
    u_ = up (rc rHCb) $ map (upB (rc rWCb)) u
    v_ = up (rc rHCr) $ map (upB (rc rWCr)) v


type PixelYUV = (Int,Int,Int)


hGetPic :: PicInfo -> Handle -> IO Picture
hGetPic i h = do
  let len = picSize i
  raw <- B.hGet h len
  if B.length raw == fromIntegral len
    then return $ packPic (i,raw)
    else return . packPic . (,) i . B.append raw
                . B.pack $ replicate (len - fromIntegral (B.length raw)) 128

getPic :: PicInfo -> FilePath -> Int -> IO (Maybe Picture)
getPic i f n = withBinaryFile f ReadMode doGet
  where
    doGet h = do
      hSeek h AbsoluteSeek $ toInteger $ n * picSize i
      eof <- hIsEOF h
      if eof
        then return Nothing
        else do
          p <- hGetPic i h
          return $ Just p

hGetPics :: PicInfo -> Handle -> IO [Picture]
hGetPics i h = do
  eof <- hIsEOF h
  if eof
    then return []
    else do
      p <- hGetPic i h
      ps <- hGetPics i h
      return $ p:ps

getPics :: PicInfo -> FilePath -> IO [Picture]
getPics i f = withBinaryFile f ReadMode (hGetPics i)
